#include "HdrProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "Render/Texture/CustomTexture.h"
#include "LocalShader.h"
#include "Render/VulkanTool.h"

void HdrProcess::OnInit()
{
	mMaxFrame = 100000;
	mFrame = 0;
	mLastView = mGameScene->GetCamera().GetView();
	mLastProj = mGameScene->GetCamera().GetProj();

	Ubo ubo;
	ubo.projview = mLastProj * mLastView;
	ubo.viewInverse = glm::inverse(mLastView);
	ubo.projInverse = glm::inverse(mLastProj);
	ubo.frame = mFrame;

	ENQUEUE_RENDER_COMMAND(BuildModel)([this, ubo](VkCommandBuffer commandBuffer)
		{
			Ubo uboData = ubo;

			mUboBuffer = CustomBuffer::Builder()
				.SetData(&uboData)
				.SetSize(sizeof(Ubo))
				.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
				.Build();

			mHdrTexture = Picture2DTexture::Builder()
				.SetFormat(VK_FORMAT_R32G32B32A32_SFLOAT)
				.SetPath("Hdr\\LoftHall_2k.hdr")
				.Build(this, "HdrTexture");

			mHdrSampler = RHISampler::Builder()
				.SetAddressModeU(VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE)
				.SetAddressModeV(VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE)
				.SetAddressModeW(VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE)
				.SetCompareOp(VK_COMPARE_OP_NEVER)
				.SetBorderColor(VK_BORDER_COLOR_FLOAT_OPAQUE_WHITE)
				.Build(this, "HdrSampler");

			std::vector<uint8> hdrCache;
			uint32 hdrWidth = 0;
			uint32 hdrHeight = 0;
			GenerateHdrCache("Hdr\\LoftHall_2k.hdr", hdrCache, hdrWidth, hdrHeight);
			mCacheTexture = CustomTexture::Builder()
				.SetWidth(hdrWidth)
				.SetHeight(hdrHeight)
				.SetFormat(VK_FORMAT_R32G32B32A32_SFLOAT)
				.SetUsage(VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetData(hdrCache.data())
				.Build();

			// Check if requested image format supports image storage operations
			VkFormatProperties formatProperties;
			vkGetPhysicalDeviceFormatProperties(GVulkanRHI->GetPhysicalDevice(), VK_FORMAT_R8G8B8A8_UNORM, &formatProperties);
			assert(formatProperties.optimalTilingFeatures & VK_FORMAT_FEATURE_STORAGE_IMAGE_BIT);

			mResultTexture = CustomTexture::Builder()
				.SetUsage(VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_STORAGE_BIT)
				.Build();

			mColorTexture = CustomTexture::Builder()
				.SetUsage(VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetMultiple(true)
				.Build(this, "ColorRenderTarget");

			TextureTransfer()
				.SetSrcStageMask(VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR)
				.SetSrcAccessMask(VK_ACCESS_MEMORY_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR)
				.SetDstAccessMask(VK_ACCESS_MEMORY_WRITE_BIT)
				.SetNewLayout(VK_IMAGE_LAYOUT_GENERAL)
				.SetTexture(mResultTexture)
				.Transfer(commandBuffer);

			TextureTransfer()
				.SetSrcStageMask(VK_PIPELINE_STAGE_TRANSFER_BIT)
				.SetSrcAccessMask(VK_ACCESS_TRANSFER_WRITE_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetNewLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.SetTexture(mColorTexture)
				.Transfer(commandBuffer, true);

			mDescriptor = RHIDescriptor::Builder()
				.PushShader<HdrVertexShader>()
				.PushShader<HdrFragmentShader>()
				.SetDescriptorCount(0, 1)
				.Build();


			mDescriptor->SetValue("ubo", mUboBuffer);
			mDescriptor->SetValue("resultImage", mResultTexture, VK_IMAGE_LAYOUT_GENERAL);
			mDescriptor->SetValue("samplerHdr", mHdrTexture, mHdrSampler);
			mDescriptor->SetValue("samplerCache", mCacheTexture, mHdrSampler);
			mDescriptorMask = mDescriptor->AllocateDescriptor();


			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.Build();

			mRenderPass = RHIRenderPass::Builder()
				.AddPipeline(mPipeline)
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.Build();

			mFrameBuffer = RHIFrameBuffer::Builder()
				.SetWidth(GRenderState->GetWidth())
				.SetHeight(GRenderState->GetHeight())
				.SetRenderPass(mRenderPass)
				.Build();
		});
}

void HdrProcess::OnTick(float deltaTime)
{
	glm::mat4 currView = mGameScene->GetCamera().GetView();
	glm::mat4 currProj = mGameScene->GetCamera().GetProj();

	if (memcmp(&mLastView, &currView, sizeof(glm::mat4)) == 0 && memcmp(&mLastProj, &currProj, sizeof(glm::mat4)) == 0)
	{
		// No change, frame ++ if less than maxframe
		if (mFrame < mMaxFrame)
		{
			mFrame++;
		}
		else
		{
			// Do not render ray until matrix change
			return;
		}
	}
	else
	{
		// Reset frame and matrix
		mFrame = 0;
		mLastView = currView;
		mLastProj = currProj;
	}


	Ubo ubo;
	ubo.projview = mLastProj * mLastView;
	ubo.viewInverse = glm::inverse(mLastView);
	ubo.projInverse = glm::inverse(mLastProj);
	ubo.frame = mFrame;

	ENQUEUE_RENDER_COMMAND(RecordCommand)([this, ubo](VkCommandBuffer commandBuffer)
		{
			Ubo uboData = ubo;

			mUboBuffer->WriteMemory(&uboData);

			VkRenderPassBeginInfo renderPassBeginInfo = mFrameBuffer->GetRenderPassBeginInfo();

			vkCmdBeginRenderPass(commandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

			vkCmdSetViewport(commandBuffer, 0, 1, &GRenderState->GetViewport());
			vkCmdSetScissor(commandBuffer, 0, 1, &GRenderState->GetScissor());

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());

			vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipelineLayout(), 0, 1, &mDescriptorMask.GetDescriptorSet(), 0, 0);

			vkCmdDraw(commandBuffer, 3, 1, 0, 0);

			vkCmdEndRenderPass(commandBuffer);

			TextureTransfer()
				.SetSrcStageMask(VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR)
				.SetSrcAccessMask(VK_ACCESS_MEMORY_WRITE_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_TRANSFER_BIT)
				.SetDstAccessMask(VK_ACCESS_TRANSFER_READ_BIT)
				.SetNewLayout(VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL)
				.SetTexture(mResultTexture)
				.Transfer(commandBuffer);

			TextureTransfer()
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_TRANSFER_BIT)
				.SetDstAccessMask(VK_ACCESS_TRANSFER_WRITE_BIT)
				.SetNewLayout(VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL)
				.SetTexture(mColorTexture)
				.Transfer(commandBuffer);

			VkImageCopy copyRegion{};
			copyRegion.srcSubresource = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1 };
			copyRegion.srcOffset = { 0, 0, 0 };
			copyRegion.dstSubresource = { VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1 };
			copyRegion.dstOffset = { 0, 0, 0 };
			copyRegion.extent = { GRenderState->GetWidth(), GRenderState->GetHeight(), 1 };
			vkCmdCopyImage(
				commandBuffer,
				mResultTexture->GetImage(),
				VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
				mColorTexture->GetImage(),
				VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
				1, &copyRegion);

			TextureTransfer()
				.SetSrcStageMask(VK_PIPELINE_STAGE_TRANSFER_BIT)
				.SetSrcAccessMask(VK_ACCESS_TRANSFER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR)
				.SetDstAccessMask(VK_ACCESS_MEMORY_WRITE_BIT)
				.SetNewLayout(VK_IMAGE_LAYOUT_GENERAL)
				.SetTexture(mResultTexture)
				.Transfer(commandBuffer);

			TextureTransfer()
				.SetSrcStageMask(VK_PIPELINE_STAGE_TRANSFER_BIT)
				.SetSrcAccessMask(VK_ACCESS_TRANSFER_WRITE_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetNewLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.SetTexture(mColorTexture)
				.Transfer(commandBuffer);

		});
}

void HdrProcess::OnResize()
{
	mResultTexture = CustomTexture::Builder()
		.SetUsage(VK_IMAGE_USAGE_TRANSFER_SRC_BIT | VK_IMAGE_USAGE_STORAGE_BIT)
		.Build();

	mColorTexture = CustomTexture::Builder()
		.SetUsage(VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetMultiple(true)
		.Build(this, "ColorRenderTarget");

	VkCommandBuffer commandBuffer = GRenderState->BeginSingleCommand();

	TextureTransfer()
		.SetSrcStageMask(VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR)
		.SetSrcAccessMask(VK_ACCESS_MEMORY_READ_BIT)
		.SetDstStageMask(VK_PIPELINE_STAGE_RAY_TRACING_SHADER_BIT_KHR)
		.SetDstAccessMask(VK_ACCESS_MEMORY_WRITE_BIT)
		.SetNewLayout(VK_IMAGE_LAYOUT_GENERAL)
		.SetTexture(mResultTexture)
		.Transfer(commandBuffer);

	TextureTransfer()
		.SetSrcStageMask(VK_PIPELINE_STAGE_TRANSFER_BIT)
		.SetSrcAccessMask(VK_ACCESS_TRANSFER_WRITE_BIT)
		.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
		.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
		.SetNewLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.SetTexture(mColorTexture)
		.Transfer(commandBuffer, true);

	GRenderState->FlushSingleCommand(commandBuffer);

	mDescriptor->SetValue("resultImage", mResultTexture, VK_IMAGE_LAYOUT_GENERAL);
	mDescriptor->UpdateDescriptor(mDescriptorMask);

	mFrameBuffer = RHIFrameBuffer::Builder()
		.SetWidth(GRenderState->GetWidth())
		.SetHeight(GRenderState->GetHeight())
		.SetRenderPass(mRenderPass)
		.Build();
}


